import cv2
import numpy as np
from skimage import filters
from matplotlib import pyplot


def sort_line(lines):
    line = []
    line_sort = []
    for i in lines:
        line.append(i[0][1])
    line = np.sort(line)
    for i in range(len(lines)):
        for j in range(len(line)):
            if line[i] == lines[j][0][1]:
                line_sort.append(lines[j])
    return line_sort


def k_means(data):
    data = np.float32(data)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    flags = cv2.KMEANS_RANDOM_CENTERS
    compactness, labels, centers = cv2.kmeans(data, 3, None, criteria, 2, flags)
    mins_0 = []
    mins_1 = []
    mins_2 = []
    mins_3 = []
    mins_4 = []
    mins_5 = []
    for i in range(len(data)):
        if labels[i] == 0:
            mins_0.append(data[i])
        if labels[i] == 1:
            mins_1.append(data[i])
        if labels[i] == 2:
            mins_2.append(data[i])
        if labels[i] == 3:
            mins_3.append(data[i])
        if labels[i] == 4:
            mins_4.append(data[i])
        if labels[i] == 5:
            mins_5.append(data[i])
    num_max = np.max((len(mins_0), len(mins_1), len(mins_2)))
    if len(mins_0) == num_max:
        return np.mean(mins_0)
    if len(mins_1) == num_max:
        return np.mean(mins_1)
    if len(mins_2) == num_max:
        return np.mean(mins_2)


def k_means_1(data):
    data = np.float32(data)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    flags = cv2.KMEANS_RANDOM_CENTERS
    compactness, labels, centers = cv2.kmeans(data, 3, None, criteria, 2, flags)
    mins_0 = []
    mins_1 = []
    mins_2 = []
    mins_3 = []
    mins_4 = []
    mins_5 = []
    for i in range(len(data)):
        if labels[i] == 0:
            mins_0.append(data[i])
        if labels[i] == 1:
            mins_1.append(data[i])
        if labels[i] == 2:
            mins_2.append(data[i])
        if labels[i] == 3:
            mins_3.append(data[i])
        if labels[i] == 4:
            mins_4.append(data[i])
        if labels[i] == 5:
            mins_5.append(data[i])
    num_max = np.max((len(mins_0), len(mins_1), len(mins_2)))
    if len(mins_0) == num_max:
        return np.mean(mins_0)
    if len(mins_1) == num_max:
        return np.mean(mins_1)
    if len(mins_2) == num_max:
        return np.mean(mins_2)


def deal_img(img):
    img_mat = np.zeros(img.shape)
    a = img.shape[0]
    b = img.shape[1]
    c = min([a, b])
    for i in range(5):
        if int(c / 2) % 2 == 1:
            c = int(c / 2)
        else:
            c = int(c / 2) + 1
        ker = (c, c)
        img_blur = cv2.GaussianBlur(img, ker, 0.0, sigmaY=5.0)
        y = cv2.Sobel(img_blur, cv2.CV_16S, 0, 1)
        absY = cv2.convertScaleAbs(y)
        absY = area_threshold(absY)
        img_mat = img_mat + absY
    return absY


def area_threshold(gray_img):
    """
    区域二值化
    :param gray_img: 边界图
    :return: 二值图
    """
    imgs = []
    img_mat = np.zeros(gray_img.shape, np.uint8)
    for i in range(3):
        img = gray_img[int(gray_img.shape[0]/3*i):int(gray_img.shape[0]/3*(i+1)), 0:gray_img.shape[1]]
        thresh_img = filters.threshold_otsu(img)
        ret, img = cv2.threshold(img, thresh_img, 255, cv2.THRESH_BINARY)
        imgs.append(img)
    for j in range(len(imgs)):
        img_mat[int(gray_img.shape[0]/3*j):int(gray_img.shape[0]/3*(j+1)), 0:gray_img.shape[1]] = imgs[j]
    return img_mat


def draw_hist(lenths):
    pyplot.hist(lenths, 100)
    pyplot.xlabel('lenth')
    pyplot.xlim(min(lenths), max(lenths))
    pyplot.ylabel('Frequency')
    pyplot.title('Lenth Of Fake Urls')
    pyplot.show()


def get_mat(lines, img):
    """
    获得变换矩阵
    :param lines:检测的直线集合
    :param img: 原图（利用图像的长和宽)
    :return: 返回变换矩阵
    """
    point = []
    img_up = []
    for i in lines:
        if i[0][1] < img.shape[0]/3:
            theta = (i[0][1] - i[0][3])/(i[0][3] - i[0][0])
            img_up.append(theta)
    num = 0
    for j in range(len(img_up)):
        if img_up[0]*img_up[j] < 0:
            num = num+1
    lines = sort_line(lines)
    point = [(lines[0][0][0], lines[0][0][1]), (lines[0][0][2], lines[0][0][3]), (lines[len(lines)-1][0][0], lines[len(lines)-1][0][1]), (lines[len(lines)-1][0][2], lines[len(lines)-1][0][3])]
    return point


def delete_line(lines):
    img_theta = []
    line_0 = lines.copy()
    for i in lines:
        theta = (i[1] - i[3]) / (i[0] - i[2])
        img_theta.append(theta)
    if len(img_theta) > 2:
        theta = k_means(img_theta)
        cnt = 0
        for i in lines:
            for j in lines:
                if ((i[3]-j[3])*(i[1]-j[1])) < 0:
                    theta_1 = (i[1] - i[3]) / (i[0] - i[2])
                    theta_2 = (j[1] - j[3]) / (j[0] - j[2])
                    if abs(theta-theta_1) > abs(theta-theta_2):
                        if i in line_0:
                            line_0.remove(i)
                    else:
                        if j in line_0:
                            line_0.remove(j)
        return line_0
    else:
        return lines


def merge_line(lines, img):
    img_theta = []
    lines = sort_line(lines)
    for i in lines:
        theta = (i[0][3] - i[0][1]) / (i[0][2] - i[0][0])
        img_theta.append(theta)
    theta_all = np.mean(img_theta)
    thetas = []
    line = []
    line_0 = []
    for i in lines:
        if i[0][2] != i[0][0]:
            k = (i[0][3] - i[0][1]) / (i[0][2] - i[0][0])
        else:
            k = 0
        b = i[0][1] - k * i[0][0]
        line_0.append([k, b])
    for i in range(len(lines) - 1):
        if abs(line_0[i][1] - line_0[i + 1][1]) < img.shape[0] / 50:
            thetas.append(line_0[i])
        else:
            differ = []
            if len(thetas) > 0:
                for j in thetas:
                    differ.append(abs(j[0]-theta_all))
                min_dif = min(differ)
                for t in thetas:
                    if abs(t[0]-theta_all) == min_dif:
                        line.append([0, int(t[1]), img.shape[1], int(img.shape[1] * t[0] + t[1])])
                        thetas.clear()
                        break
    return line


def  getline(lines, img):
    """
    辅助函数（将得到的直线可视化）
    :param lines: 直线检测得到的直线
    :param img: 图像
    :return: 返回直线的可视化图像
    """
    correct_line = []
    lines = sort_line(lines)
    lines = merge_line(lines, img)
    lines = delete_line(lines)
    for line in lines:
        x1 = line[0]
        y1 = line[1]
        x2 = line[2]
        y2 = line[3]
        if (x2 - x1) != 0:
            theta = (y1 - y2) / (x2 - x1)
            if abs(theta) < 0.6 and y1 < img.shape[0]*9/10 and y1 > img.shape[0]/10:
                correct_line.append([[x1, y1, x2, y2]])
    return correct_line


def Transform(img, points_1):
    """
    进行透射变换
    :param img:图像
    :param points_1:变换矩阵
    :return: 透射变换后图像
    """
    rows, cols, ch = img.shape
    pts1 = np.float32(points_1)
    pts2 = np.float32([[0, 0], [int(cols), 0], [0, int(max([points_1[2][1]-points_1[0][1], points_1[3][1]-points_1[1][1]]))], [int(cols), int(max([points_1[2][1]-points_1[0][1], points_1[3][1]-points_1[1][1]]))]])
    M = cv2.getPerspectiveTransform(pts1, pts2)
    dst = cv2.warpPerspective(img, M, (int(img.shape[1]), int(max([points_1[2][1]-points_1[0][1], points_1[3][1]-points_1[1][1]]))), flags=cv2.INTER_LINEAR)
    M_1 = cv2.getPerspectiveTransform(pts2, pts1)


def find_ker(img):
    h, w = img.shape
    min_x = np.min([h, w])
    min_x = int(min_x / 50 )
    if min_x % 2 == 0:
        min_x = min_x - 1
    return (min_x, min_x)


def img_resize(img):
    h = img.shape[0]
    w = img.shape[1]
    max_x = max([h, w])
    if max_x == h:
        h = 600
        w = w * 600 / h
    if max_x == w:
        w = 600
        h = h * 600 / w
    img = cv2.resize(img, (int(h), int(w)))
    return img


def getPoint(line, img):
    x1 = line[0][0]
    y1 = line[0][1]
    x2 = line[0][2]
    y2 = line[0][3]
    if x2 != x1:
        k = (y2-y1)/(x2-x1)
    else:
        k = 0
    b = y1 - k*x1
    point_1 = [0, b]
    if k != 0:
        point_2 = [img.shape[1], img.shape[1]*k+b]
    else:
        point_2 = [img.shape[1], b]
    return [point_1, point_2]


def get_mask(line_point, img):
    img_mask = np.ones((img.shape[0], img.shape[1], 3), np.uint8)
    line_point = [[line_point[0], line_point[1], line_point[3], line_point[2]]]
    line_point = np.array(line_point)
    cv2.fillPoly(img_mask, line_point, 0)
    img_mask = img*img_mask
    return img_mask


def i_correct(img):
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)
    l, a, b = cv2.split(img_gray)
    img_kea = find_ker(l)
    img_gray = cv2.GaussianBlur(l, img_kea, 5.0)
    img_gray = cv2.GaussianBlur(img_gray, img_kea, 5.0)
    absY = deal_img(l)
    bw = np.zeros(absY.shape, np.uint8)
    ret, contours, hierarchy = cv2.findContours(absY, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    for i in contours:
        if cv2.contourArea(i) > img.shape[1]:
            i = [i]
            cv2.drawContours(bw, i, -1, (255, 255, 255), -1)
    bw = cv2.dilate(bw, cv2.getStructuringElement(cv2.MORPH_RECT, (2, 3)))
    lines = cv2.HoughLinesP(bw, 1, np.pi / 180, int(img_gray.shape[1] / 5),
                            minLineLength=int(img_gray.shape[1] / 5), maxLineGap=int(img_gray.shape[1] / 500))
    correct_l = getline(lines, img)
    if len(correct_l) > 3:
        point_1 = get_mat(correct_l, img)
        dst, M_1 = Transform(img, point_1)
        dst_mask = get_mask(point_1, img)
        cv2.imshow('mask', dst_mask)
        result = {'1': dst, '2': M_1, '3': dst_mask}
        return result
    else:
        dst = img
        result = {'1': dst}
        return result
